from MySAM2 import SAM2
import numpy as np
import cv2
import torch
import torchvision.transforms.functional as F
from PIL import Image, ImageDraw

if __name__ == '__main__':
    model = SAM2(r'configs/sam2.1_hiera_l.yaml', r'weights/sam2.1_hiera_large.pt')
    img_np = cv2.imread(r'testSAM2.jpg')
    img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
    out = model.forward(image=img_np, points=None, labels=None, box=np.array([[212, 52, 446, 132]]), mask_logit=None)

    img_pil = Image.fromarray(np.uint8(img_np))
    img_draw = ImageDraw.Draw(img_pil)
    img_draw.rectangle([212, 52, 446, 132],outline='red',width=3)

    img_pil.show()
    F.to_pil_image(torch.where(out[0] >= 0.5, 1.0, 0.0)).show()
